import json
import os
from modify_json_images_videos import modify_json_images_videos
from replace_prefix_keys import replace_prefix_keys
def find_max_id(data):
    """
    从给定的JSON结构中找出最大编号的值
    
    参数:
        data (dict): 包含节点信息的字典
        
    返回:
        int: 最大编号的值
    """
    max_id = 0
    for key in data.keys():
        if key.isdigit():  # 检查键是否为数字
            current_id = int(key)
            if current_id > max_id:
                max_id = current_id
    return max_id

def update_node_numbers(data, max_value):
    # 创建一个新的字典来存储更新后的节点
    new_data = {}
    
    # 遍历原始字典中的每个节点
    for old_key, node in data.items():
        # 将键转换为整数并加上max_value
        new_key = str(int(old_key) + max_value)
        
        # 创建节点的深拷贝以避免修改原始数据
        new_node = node.copy()
        
        # 检查节点是否有inputs字段
        if 'inputs' in new_node:
            # 遍历inputs中的每个值
            for input_key, input_value in new_node['inputs'].items():
                # 如果值是列表且第一个元素是字符串（可能是节点引用）
                if isinstance(input_value, list) and len(input_value) > 0 and isinstance(input_value[0], str):
                    try:
                        # 尝试将第一个元素转换为整数并加上max_value
                        new_ref = str(int(input_value[0]) + max_value)
                        new_node['inputs'][input_key] = [new_ref] + input_value[1:]
                    except ValueError:
                        # 如果不能转换为整数，保持原样
                        pass
        
        # 将更新后的节点添加到新字典中
        new_data[new_key] = new_node
    
    return new_data
def find_highest_node_by_type(workflow_dict, class_type):
    matching_nodes = []
    
    for node_id, node_data in workflow_dict.items():
        if node_data.get('class_type') == class_type:
            # Convert node_id to integer for comparison
            try:
                node_id_int = int(node_id)
                matching_nodes.append(node_id_int)
            except ValueError:
                # Handle case where node_id isn't a number (though unlikely in this context)
                pass
    
    if not matching_nodes:
        return None
    
    return str(max(matching_nodes))
def modify_node_input(data, class_type, input_class, num):
    """
    修改指定节点的输入类对应的连接编号
    
    参数:
        data: 包含所有节点数据的字典
        class_type: 要修改的节点的class_type
        input_class: 要修改的输入字段名
        num: 新的连接编号(字符串)
        
    返回:
        修改后的data字典
    """
    # 遍历所有节点
    for node_id, node_data in data.items():
        # 检查是否是目标节点类型
        if node_data.get("class_type") == class_type:
            # 检查是否有inputs字段和指定的input_class
            if "inputs" in node_data and input_class in node_data["inputs"]:
                # 获取当前输入值
                current_value = node_data["inputs"][input_class]
                
                # 如果是列表形式且长度为2且第二个元素是数字
                if isinstance(current_value, list) and len(current_value) == 2 and isinstance(current_value[1], int):
                    # 修改第一个元素为新的num
                    node_data["inputs"][input_class][0] = str(num)
                    
    return data


def merge_dicts_update(dict1, dict2):
    """
    使用update方法合并两个字典
    注意：这会修改dict1的内容
    """
    result = dict1.copy()  # 避免修改原始字典
    result.update(dict2)
    return result

def merge_2_workflow(data1,data2,data1_output_id,):
    """
    对于有依赖关系的两个框架A和B，
    将A中的资源输出节点连向B中的资源输入节点。
    
    
    """

    # 找出第一个工作流的最大节点

    max_id = find_max_id(data1)
    
    # 找出第一个工作流的输出节点的对应编号
    
    # id = find_highest_node_by_type(data1,"VAEDecode")

    # 将第二个工作流的全部节点加上偏移值,使得不重复

    data2 = update_node_numbers(data2,max_value=max_id)

    # 找出第二个工作流符合图像输入的节点，指向第一个工作流的输出
    
    

    data2 = modify_node_input(data2,"CLIPVisionEncode","image",id)
    data2 = modify_node_input(data2,"WanImageToVideo","start_image",id)
    data2 = modify_node_input(data2,"VAEEncode","pixels",id)
    # 最后合并工作流

    
    
    last = merge_dicts_update(data1,data2)


    # 写入JSON文件（带缩进格式化）
    with open('generated_workflow/dad.json', 'w', encoding='utf-8') as f:
        json.dump(last, f, ensure_ascii=False, indent=4)
    
    return last

    pass

def modify_all_positive_texts(json_data, new_text):
    """
    修改JSON数据中所有positive节点对应的文本输入框的text值
    
    参数:
        json_data: 包含工作流数据的JSON对象
        new_text: 要替换的新文本
    
    返回:
        修改后的JSON数据
    """
    # 遍历所有节点
    for node_id, node_data in json_data.items():
        # 检查节点是否有inputs
        if "inputs" not in node_data:
            continue
            
        inputs = node_data["inputs"]
        
        # 检查是否有positive输入（可能是列表或字典）
        if "positive" in inputs:
            positive_ref = inputs["positive"]
            
            # 如果positive是列表（如 ["6", 0]）
            if isinstance(positive_ref, list) and len(positive_ref) >= 1:
                ref_node_id = str(positive_ref[0])  # 确保是字符串
                
                # 找到对应的节点
                if ref_node_id in json_data:
                    ref_node = json_data[ref_node_id]
                    
                    # 检查该节点是否有text输入（不限定CLIPTextEncode）
                    if "inputs" in ref_node and "text" in ref_node["inputs"]:
                        ref_node["inputs"]["text"] = new_text
                        
            # 如果positive是字典（某些工作流可能不同）
            elif isinstance(positive_ref, dict) and "text" in positive_ref:
                positive_ref["text"] = new_text
    
    return json_data
# 定义 load_workflow 函数
def load_workflow(file_path):
    """从 JSON 文件加载工作流数据"""
    with open(file_path, 'r', encoding='utf-8') as f:
        return json.load(f)

def find_start_end_nodes(dag):
    # 收集所有节点ID
    all_nodes = set(dag.keys())
    
    # 构建依赖关系：记录每个节点的输入来源
    dependencies = {node_id: set() for node_id in all_nodes}
    # 构建被依赖关系：记录每个节点的输出去向
    dependents = {node_id: set() for node_id in all_nodes}
    
    for node_id, node_data in dag.items():
        inputs = node_data.get('inputs', {})
        
        # 遍历所有输入连接
        for input_name, input_value in inputs.items():
            # 输入连接可能是列表形式 [source_node_id, output_index]
            if isinstance(input_value, list) and len(input_value) >= 2:
                source_node = str(input_value[0])  # 转换为字符串以确保类型一致
                if source_node in all_nodes:  # 确保引用的节点存在
                    dependencies[node_id].add(source_node)
                    dependents[source_node].add(node_id)
    
    # 起始节点：没有依赖其他节点的节点
    start_nodes = [node_id for node_id in all_nodes if not dependencies[node_id]]
    
    # 结束节点：没有被任何其他节点依赖的节点
    end_nodes = [node_id for node_id in all_nodes if not dependents[node_id]]
    
    return start_nodes, end_nodes

def find_direct_dependents(dag, node_id):
    """
    找出所有直接依赖指定节点的节点ID列表（不递归查找间接依赖）
    
    参数:
        dag: DAG结构，字典格式
        node_id: 要查找的节点ID（字符串或数字）
        
    返回:
        直接依赖该节点的所有节点ID列表（按输入顺序排序）
    """
    node_id = str(node_id)  # 确保统一为字符串格式
    dependents = []
    
    for current_node_id, node_data in dag.items():
        inputs = node_data.get('inputs', {})
        
        # 检查所有输入连接
        for input_value in inputs.values():
            # 处理单个连接或连接列表
            connections = [input_value] if not isinstance(input_value, list) else input_value
            
            for conn in connections:
                # 检查是否是节点连接格式 (如 ["50", 0] 或 "50")
                if isinstance(conn, list) and len(conn) > 0:
                    dep_id = str(conn[0]) if len(conn) > 0 else None
                elif isinstance(conn, str):
                    dep_id = conn
                else:
                    continue  # 忽略非连接格式的输入
                
                if dep_id == node_id:
                    dependents.append(current_node_id)
                    break  # 避免重复记录同一个节点
    
    return dependents

def find_min_text_or_prompt_node_id(dag_data):
    """
    返回拥有text或prompt字段的最小节点ID
    """
    nodes = []
    
    # 遍历所有节点
    for node_id, node_data in dag_data.items():
        inputs = node_data.get("inputs", {})
        # 检查节点是否有inputs且包含text或prompt字段
        if "text" in inputs or "prompt" in inputs:
            nodes.append(node_id)
    
    # 如果没有找到任何符合条件的节点，返回空字符串
    if not nodes:
        return ""
    
    # 找到ID值最小的节点（转换为整数比较）
    min_id = min(nodes, key=lambda x: int(x))
    
    return min_id


def find_dependent_nodes(data):
    dependent_nodes = set()
    
    for key, value in data.items():
        if isinstance(value, list):
            for item in value:
                if isinstance(item, str) and item.isdigit():
                    dependent_nodes.add(item)
        # 不处理非字符串的数字（如 0, 1, 2）
    
    return sorted(dependent_nodes)

def find_matching_field(input_dict):
    """
    遍历输入字典中的inputs字段，返回第一个在预设列表中存在的字段名。
    硬编码的检查字段列表: ["vae", "pixels", "other_field"]
    
    参数:
        input_dict (dict): 包含inputs字段的字典
        
    返回:
        str: 第一个匹配的字段名，如果没有匹配则返回None
    """
    fields_to_check = ["image","pixels","images","start_image","init_image","video"]  # 硬编码的字段列表
    inputs = input_dict.get('inputs', {})
    for field in fields_to_check:
        if field in inputs:
            return field
    return None
def merge_two_flow(data1,data2):
    # 找出第一个工作流的最大节点

    max_id = find_max_id(data1)

    # 将第二个工作流的全部节点加上偏移值,使得不重复

    data2 = update_node_numbers(data2,max_value=max_id)

    ## 找出data1的所有结束节点

    _,data1_last_ids = find_start_end_nodes(data1)
    print("全部输出:")
    print(data1_last_ids)
    ## 按照data1的结束节点中的input找出最近的输出节点

    data1_last_output_id = find_dependent_nodes(data1[data1_last_ids[0]]["inputs"])[0]
    print(data1_last_output_id)
    
    ## 找出data2所有的开始节点

    data2_start_ids,_ = find_start_end_nodes(data2)
    print(data2_start_ids)

    ## 过滤data2的开始节点，只保留与资源输入有关的

    source_input_ids = [] ##候选的资源输入节点

    for id in data2_start_ids:
        if data2[id]['class_type'] in ["LoadImage"]:
            source_input_ids.append(id)
    print(source_input_ids)
    ## 找出data2开始节点的依赖节点

    data2_input_ids = []
    for id in source_input_ids:
        print("次级输入节点")
        data2_input_ids = find_direct_dependents(data2,id)
        print(data2_input_ids)

    ## 将data1的输出节点连向data2的次级输入节点

    for id in data2_input_ids:
        print(id)
        match_filed = find_matching_field(data2[id])
        print(match_filed)
        data2[id]["inputs"][match_filed][0] = data1_last_output_id

        pass
    ## 返回第二个工作流

    return data2
def load_task_paths(json_file_path):
    try:
        with open(json_file_path, 'r', encoding='utf-8') as file:
            task_paths = json.load(file)
        return task_paths
    except FileNotFoundError:
        print(f"错误：找不到配置文件 {json_file_path}")
        return {}
    except json.JSONDecodeError:
        print(f"错误：配置文件 {json_file_path} 格式不正确")
        return {}
def run(task_name,resource,input):
    config_file = 'tasks_config.json'  # 如果文件不在同一目录，请提供完整路径

    # 加载任务路径
    task_basic_path = load_task_paths(config_file)
    for i in range(len(input)):
        input[i].append(load_workflow(task_basic_path[input[i][0]]))

    

    for i in range(len(input)):
        # 针对每个单独的子任务基本流程进行提示词构造并赋值
        id = find_min_text_or_prompt_node_id(input[i][-1])
        if id != "":
            # 获取text或prompt的值
            value = input[i][1]
            
            # 检查是否存在text键
            if "text" in input[i][-1][id]["inputs"]:
                input[i][-1][id]["inputs"]["text"] = value
            # 检查是否存在prompt键
            elif "prompt" in input[i][-1][id]["inputs"]:
                input[i][-1][id]["inputs"]["prompt"] = value
    
    ##资源输入和目标文件名输入
    for i in range(len(input)):
        if resource != "" and resource != None:

            input[i][-1] = modify_json_images_videos(input[i][-1],resource)
        
        input[i][-1] = replace_prefix_keys(input[i][-1],task_name)
        pass
    if len(input) == 1:
        last = input[0][-1]
        ## 保存为json文件
        with open('generated_workflow/dad.json', 'w', encoding='utf-8') as f:
            json.dump(last, f, ensure_ascii=False, indent=4)
        
        os.system("python comfyui_exe.py")
        pass
    else:
        ## 对input中每一个流程增加id偏移值
        for i in range(len(input)):
            if i == 0:
                continue
            input[i][-1] = update_node_numbers(input[i][-1],find_max_id(input[i-1][-1]))
            pass
        ## 按照input中每个流程的依赖关系更新流程,其中倒数第二个索引是依赖的索引，-1是没有前置依赖，使用merge_two_flow函数
        for i in range(len(input)):
            if input[i][-2] == -1:
                continue
            input[i][-1] = merge_two_flow(input[input[i][-2]][-1],input[i][-1])



        ## 将所有流程文件，合并为一个，就是合并所有字典，不需要merge_two_flow函数,只用内置的字典合并函数就好
        last = {}
        for i in range(len(input)):
            last = merge_dicts_update(last,input[i][-1])
        ## 保存为json文件
        with open('generated_workflow/dad.json', 'w', encoding='utf-8') as f:
            json.dump(last, f, ensure_ascii=False, indent=4)
        
        os.system("python comfyui_exe.py")


    pass
if __name__ == "__main__":

    
    input = [  
        ["文生图", "A cute orange tabby cat, sitting on a white background, detailed fur, bright eyes, photorealistic style, 8K resolution, soft lighting, high detail.", -1],  
        ["文生图", "A small black dog riding a bicycle on a sunny road, cartoon-style, vibrant colors, dynamic pose, playful expression, background with trees and blue sky.", -1]]

    run(input)
